#!/usr/bin/python

import re
import nntplib as nntplibOriginal
import progressbar

class NNTPError(nntplibOriginal.NNTPError):
    def __init__(self):
        nntplibOriginal.NNTPError.init(self)

class NNTPError(nntplibOriginal.NNTPError):
    def __init__(self):
        nntplibOriginal.NNTPError.init(self)

class NNTP(nntplibOriginal.NNTP):

    start = 0
    end = 0
    nb_articles = 0

    def __init__(self, host, port=119, user=None, password=None,readermode=None):
        nntplibOriginal.NNTP.__init__(self, host, port, user, password, readermode)

    def getlongresp(self, file=None):
        """Internal: get a response plus following text from the server.
        Raise various errors if the response indicates an error."""

        use_bar=0

        if (self.start or self.end or self.nb_articles):
            #widgets = ['Downloading: ', progressbar.Percentage(), ' ', progressbar.Bar(marker='>',left='[',right=']'), ' ', progressbar.ETA(), ' ', progressbar.FileTransferSpeed()]
            widgets = ['Downloading: ', progressbar.Percentage(), ' ', progressbar.Bar(marker='>',left='[',right=']'), ' ', progressbar.ETA(), ' ','']
            
            pbar = progressbar.ProgressBar(widgets=widgets, maxval=self.nb_articles).start()
            use_bar=1        

        openedFile = None
        try:
            # If a string was passed then open a file with that name
            if isinstance(file, str):
                openedFile = file = open(file, "w")

            resp = self.getresp()

            if resp[:3] not in nntplibOriginal.LONGRESP:
                raise NNTPReplyError(resp)
            list = []
            a = 0
            while 1:
                line = self.getline()
                if use_bar:
                    if (a<=self.nb_articles):
                        pbar.update(a)
                        a += 1
                if line == '.':
                    break
                if line[:2] == '..':
                    line = line[1:]
                if file:
                    file.write(line + "\n")
                else:
                    list.append(line)
        finally:
            # If this method created the file, then it must close it
            if openedFile:
                openedFile.close()

        if use_bar:
            pbar.update(self.nb_articles)

        return resp, list


    def xover(self, str, file=None):
        """Process an XOVER command (optional server extension).  Arguments:
        - str: an article nr, a message id, or a range nr1-nr2
        Returns:
        - resp: server response if successful
        - list: list of (nr, value) strings"""

        s=str.split('-')
        self.start=int(s[0])
        self.end=int(s[1])
        self.nb_articles=self.end-self.start

        pat = re.compile('^([0-9]+) ?(.*)\n?')
        resp, lines = self.longcmd('XOVER ' + str, file)

        #for i in range(len(lines)):
        #    line = lines[i]
        #    m = pat.match(line)
        #    if m:
        #        lines[i] = m.group(1, 2)
        return resp, lines

    def xhdr(self, hdr, str, file=None):
        """Process an XHDR command (optional server extension).  Arguments:
        - hdr: the header type (e.g. 'subject')
        - str: an article nr, a message id, or a range nr1-nr2
        Returns:
        - resp: server response if successful
        - list: list of (nr, value) strings"""

        s=str.split('-')
        self.start=int(s[0])
        self.end=int(s[1])
        self.nb_articles=self.end-self.start

        pat = re.compile('^([0-9]+) ?(.*)\n?')
        resp, lines = self.longcmd('XHDR ' + hdr + ' ' + str, file)

        for i in range(len(lines)):
            line = lines[i]
            m = pat.match(line)
            if m:
                lines[i] = m.group(1, 2)
        return resp, lines

    def xpat(self, hdr, str, search, file=None):
        """Process an XPAT command (optional server extension).  Arguments:
        - hdr: the header type (e.g. 'subject')
        - str: an article nr, a message id, or a range nr1-nr2
        - search: a pattern..
        Returns:
        - resp: server response if successful
        - list: list of (nr, value) strings"""

        s=str.split('-')
        self.start=int(s[0])
        self.end=int(s[1])
        self.nb_articles=self.end-self.start

        pat = re.compile('^([0-9]+) ?(.*)\n?')
        resp, lines = self.longcmd('XPAT ' + hdr + ' ' + str + ' ' + search, file)

        for i in range(len(lines)):
            line = lines[i]
            m = pat.match(line)
            if m:
                lines[i] = m.group(1, 2)
        return resp, lines

    def overview_format(self, file=None):
        """Process a LIST command.  Return:
        - resp: server response if successful
        - list: list of (group, last, first, flag) (strings)"""

        resp, list = self.longcmd('LIST OVERVIEW.FMT', file)
        for i in range(len(list)):
            # Parse lines into "group last first flag"
            list[i] = list[i].split(':')[0]
        return resp, list

if __name__ == '__main__':
    import os
    newshost = 'news' and os.environ["NNTPSERVER"]
    if newshost.find('.') == -1:
        mode = 'readermode'
    else:
        mode = None
    s = NNTP(newshost, readermode=mode)
    resp, count, first, last, name = s.group('comp.lang.python')
    print resp
    print 'Group', name, 'has', count, 'articles, range', first, 'to', last
    resp, subs = s.xhdr('subject', first + '-' + last)
    print resp
    for item in subs:
        print "%7s %s" % item
    resp = s.quit()
    print resp

